# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

import torch
import pytest
from loguru import logger
import ttnn
import math
from tests.tt_eager.python_api_testing.sweep_tests.comparison_funcs import comp_pcc


def run_all_reduce_test(
    mesh_device,
    num_devices,
    per_chip_output_shape,
    num_links,
    math_op,
    input_dtype,
    layout,
    mem_config,
    function_level_defaults,
    num_iters=1,
    topology=ttnn.Topology.Linear,
):
    if len(mesh_device.get_device_ids()) < num_devices:
        pytest.skip(
            f"Not enough devices on machine to implement test case. Wanted {num_devices} but found {len(mesh_device.get_device_ids())}"
        )

    ttnn.synchronize_device(mesh_device)

    sub_device_stall_group = []
    compute_grid_size = mesh_device.compute_with_storage_grid_size()
    ccl_sub_device_crs = ttnn.CoreRangeSet(
        {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(compute_grid_size.x - 1, compute_grid_size.y - 1))}
    )
    worker_sub_device = ttnn.SubDevice([ccl_sub_device_crs])
    worker_sub_device_id = ttnn.SubDeviceId(0)
    sub_device_stall_group = [worker_sub_device_id]
    sub_device_manager = mesh_device.create_sub_device_manager([worker_sub_device], 0)
    mesh_device.load_sub_device_manager(sub_device_manager)
    mesh_device.set_sub_device_stall_group(sub_device_stall_group)
    # create global semaphore handles
    rs_global_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(3)]
    ag_global_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(2)]
    barrier_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(2)]

    debug = False

    logger.info(f"Per chip output shape: {per_chip_output_shape}, devices: {num_devices}")
    # Generate input tensors

    canonical_input_tensors = []
    input_tensors = []

    numel = math.prod(per_chip_output_shape)
    if debug:
        input_tensors[-1] = torch.arange(numel).reshape(per_chip_output_shape).bfloat16()
    for i in range(num_devices):
        input_tensor = torch.rand(per_chip_output_shape).bfloat16()
        canonical_input_tensors.append(input_tensor)
        input_tensor = input_tensor.view(1, -1, input_tensor.shape[2], input_tensor.shape[3])
        input_tensors.append(input_tensor)

    unchunked_input_tensor = torch.cat(input_tensors)

    assert len(canonical_input_tensors) == num_devices
    input_tensor_mesh = ttnn.from_torch(
        torch.cat(canonical_input_tensors),
        dtype=input_dtype,
        layout=layout,
        device=mesh_device,
        memory_config=mem_config,
        mesh_mapper=ttnn.create_mesh_mapper(
            mesh_device,
            ttnn.MeshMapperConfig([ttnn.PlacementReplicate(), ttnn.PlacementShard(0)], ttnn.MeshShape(1, num_devices)),
        ),
    )
    # Run the op
    for i in range(num_iters):
        output_tensor_mesh = ttnn.experimental.all_reduce_async(
            input_tensor_mesh,
            num_devices=num_devices,
            barrier_semaphores=barrier_semaphores,
            rs_global_semaphores=rs_global_semaphores,
            ag_global_semaphores=ag_global_semaphores,
            math_op=math_op,
            num_links=num_links,
            memory_config=mem_config,
            topology=ttnn.Topology.Linear,
            subdevice_id=worker_sub_device_id,
        )
        ttnn.synchronize_device(mesh_device, sub_device_ids=sub_device_stall_group)
    ttnn.synchronize_device(mesh_device, sub_device_ids=sub_device_stall_group)

    mesh_device.reset_sub_device_stall_group()

    tt_out_tensors = ttnn.get_device_tensors(output_tensor_mesh)
    logger.info(f"Compare")
    golden_canonical_out_tensor = torch.sum(unchunked_input_tensor, 0, keepdim=True)
    golden_canonical_out_tensor = golden_canonical_out_tensor.view(per_chip_output_shape)
    # Compare
    mismatch = False
    for i, t in enumerate(tt_out_tensors):
        tt_output_tensor = ttnn.to_torch(t)

        eq, output = comp_pcc(tt_output_tensor, golden_canonical_out_tensor)
        mismatch = mismatch or not eq
        if not eq:
            logger.error(f"output mismatch for tensor {i}. Mesh device ID: {mesh_device.get_device_ids()[i]}")
            if debug:
                for w in range(tt_output_tensor.shape[0]):
                    for z in range(tt_output_tensor.shape[1]):
                        for y in range(tt_output_tensor.shape[2]):
                            for x in range(tt_output_tensor.shape[3]):
                                if tt_output_tensor[w, z, y, x] != golden_canonical_out_tensor[w, z, y, x]:
                                    logger.error(
                                        f"mismatch at {w}, {z}, {y}, {x}: {tt_output_tensor[w, z, y, x]} != {golden_canonical_out_tensor[w, z, y, x]}"
                                    )

        else:
            logger.info(f"output match for tensor {i}")
    assert not mismatch, f"{i} FAILED: {output}"


def run_all_reduce_with_mesh_tensor_along_row(
    mesh_device,
    num_devices_per_line,
    per_chip_output_shape,
    num_links,
    math_op,
    input_dtype,
    layout,
    buffer_type: ttnn.BufferType,
    function_level_defaults,
    memory_config=None,
    num_all_reduce_instances: int = 1,
    num_iters: int = 1,
    cluster_axis=None,
    use_semaphore_free_all_reduce_impl: bool = False,
):
    mem_config = memory_config or ttnn.MemoryConfig(buffer_type=buffer_type)

    ttnn.synchronize_device(mesh_device)

    sub_device_stall_group = []
    compute_grid_size = mesh_device.compute_with_storage_grid_size()
    ccl_sub_device_crs = ttnn.CoreRangeSet(
        {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(compute_grid_size.x - 1, compute_grid_size.y - 1))}
    )
    worker_sub_device = ttnn.SubDevice([ccl_sub_device_crs])
    worker_sub_device_id = ttnn.SubDeviceId(0)
    sub_device_stall_group = [worker_sub_device_id]
    sub_device_manager = mesh_device.create_sub_device_manager([worker_sub_device], 0)
    mesh_device.load_sub_device_manager(sub_device_manager)
    mesh_device.set_sub_device_stall_group(sub_device_stall_group)
    # create global semaphore handles
    rs_global_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(3)]
    ag_global_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(2)]
    barrier_semaphores = [ttnn.create_global_semaphore(mesh_device, ccl_sub_device_crs, 0) for _ in range(2)]
    try:
        debug = False

        logger.info(f"Per chip output shape: {per_chip_output_shape}, devices: {num_devices_per_line}")
        # Generate input tensors
        input_tensors = []
        tt_input_tensors = []

        numel = math.prod(per_chip_output_shape)
        if debug:
            input_tensors[-1] = torch.arange(numel).reshape(per_chip_output_shape).bfloat16()
        for i in range(num_devices_per_line):
            input_tensor = torch.rand(per_chip_output_shape).bfloat16()
            tt_input_tensors.append(input_tensor.clone())
            input_tensor = input_tensor.view(1, -1, input_tensor.shape[2], input_tensor.shape[3])
            input_tensors.append(input_tensor)

        unchunked_input_tensor = torch.cat(input_tensors)

        shard_dims = (0, 1) if cluster_axis == 0 else (1, 0)
        mesh_shape = (
            (num_devices_per_line, num_all_reduce_instances)
            if cluster_axis == 0
            else (num_all_reduce_instances, num_devices_per_line)
        )
        full_input_tensor_unfractured = torch.cat(tt_input_tensors)
        full_input_tensor_unfractured = [full_input_tensor_unfractured for _ in range(num_all_reduce_instances)]
        full_input_tensor_unfractured = torch.cat(full_input_tensor_unfractured, dim=1)

        ttnn_tensor = ttnn.from_torch(
            full_input_tensor_unfractured,
            dtype=input_dtype,
            device=mesh_device,
            layout=layout,
            memory_config=mem_config,
            mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, mesh_shape=mesh_shape, dims=shard_dims),
        )
        input_tensor_mesh = ttnn.to_device(ttnn_tensor, mesh_device)

        # Run the op
        for i in range(num_iters):
            if use_semaphore_free_all_reduce_impl:
                logger.info("Using semaphore-free all-reduce implementation")
                output_tensor_mesh = ttnn.all_reduce(
                    input_tensor_mesh,
                    cluster_axis=cluster_axis,
                    subdevice_id=worker_sub_device_id,
                    num_links=num_links,
                    memory_config=mem_config,
                    topology=ttnn.Topology.Linear,
                )
            else:
                logger.info("Using experimental all-reduce implementation")
                output_tensor_mesh = ttnn.experimental.all_reduce_async(
                    input_tensor_mesh,
                    cluster_axis=cluster_axis,
                    mesh_device=mesh_device,
                    barrier_semaphores=barrier_semaphores,
                    rs_global_semaphores=rs_global_semaphores,
                    ag_global_semaphores=ag_global_semaphores,
                    math_op=math_op,
                    num_links=num_links,
                    memory_config=mem_config,
                    topology=ttnn.Topology.Linear,
                    subdevice_id=worker_sub_device_id,
                )
            ttnn.synchronize_device(mesh_device, sub_device_ids=sub_device_stall_group)
        ttnn.synchronize_device(mesh_device, sub_device_ids=sub_device_stall_group)
    except Exception as e:
        raise e
    finally:
        mesh_device.reset_sub_device_stall_group()

    tt_out_tensors = ttnn.get_device_tensors(output_tensor_mesh)

    logger.info(f"Compare")
    golden_canonical_out_tensor = torch.sum(unchunked_input_tensor, 0, keepdim=True)
    golden_canonical_out_tensor = golden_canonical_out_tensor.view(per_chip_output_shape)

    # Compare
    mismatch = False
    for i, t in enumerate(tt_out_tensors):
        tt_output_tensor = ttnn.to_torch(t)
        eq, output = comp_pcc(tt_output_tensor, golden_canonical_out_tensor)
        mismatch = mismatch or not eq
        if not eq:
            logger.error(f"output mismatch for tensor {i}. Mesh device ID: {mesh_device.get_device_ids()[i]}")
            if debug:
                for w in range(tt_output_tensor.shape[0]):
                    for z in range(tt_output_tensor.shape[1]):
                        for y in range(tt_output_tensor.shape[2]):
                            for x in range(tt_output_tensor.shape[3]):
                                if tt_output_tensor[w, z, y, x] != golden_canonical_out_tensor[w, z, y, x]:
                                    logger.error(
                                        f"mismatch at {w}, {z}, {y}, {x}: {tt_output_tensor[w, z, y, x]} != {golden_canonical_out_tensor[w, z, y, x]}"
                                    )

        else:
            logger.info(f"output match for tensor {i}")
    assert not mismatch, f"{i} FAILED: {output}"


# Enumerate the post-commit cases explicitly
@pytest.mark.parametrize(
    "num_devices, num_links, per_chip_output_shape, layout",
    [
        (4, 2, [1, 4, 32, 2304], ttnn.TILE_LAYOUT),
        (4, 2, [4, 1, 64, 1024], ttnn.TILE_LAYOUT),
        (4, 2, [3, 2, 90, 2040], ttnn.TILE_LAYOUT),
        (4, 2, [16, 1, 16, 512], ttnn.ROW_MAJOR_LAYOUT),
        (4, 2, [1, 1, 250, 2048], ttnn.ROW_MAJOR_LAYOUT),
        (4, 2, [2, 2, 350, 350], ttnn.ROW_MAJOR_LAYOUT),
    ],
)
@pytest.mark.parametrize(
    "input_dtype",
    [
        ttnn.bfloat16,
    ],
)
@pytest.mark.parametrize(
    "buffer_type",
    [
        ttnn.BufferType.DRAM,
        ttnn.BufferType.L1,
    ],
)
@pytest.mark.parametrize("replication_factor", [8])
@pytest.mark.parametrize("mesh_device", [pytest.param((8, 4), id="8x4_grid")], indirect=True)
@pytest.mark.parametrize("math_op", [ttnn.ReduceType.Sum])
@pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_line_all_reduce_on_TG_rows_post_commit(
    mesh_device,
    num_devices,
    per_chip_output_shape,
    num_links,
    math_op,
    input_dtype,
    layout,
    buffer_type,
    function_level_defaults,
    replication_factor,
    num_iters=16,
):
    if mesh_device.get_num_devices() != 32:
        pytest.skip("Not TG!")

    run_all_reduce_with_mesh_tensor_along_row(
        mesh_device,
        num_devices,
        per_chip_output_shape,
        num_links,
        math_op,
        input_dtype,
        layout,
        buffer_type,
        function_level_defaults,
        num_iters=num_iters,
        num_all_reduce_instances=replication_factor,
        cluster_axis=1,
    )


@pytest.mark.parametrize(
    "num_devices, num_links, per_chip_output_shape, layout",
    [
        (8, 1, [1, 8, 32, 1280], ttnn.TILE_LAYOUT),
    ],
)
@pytest.mark.parametrize(
    "input_dtype",
    [
        ttnn.bfloat16,
    ],
)
@pytest.mark.parametrize(
    "buffer_type",
    [
        ttnn.BufferType.DRAM,
    ],
)
@pytest.mark.parametrize("replication_factor", [4])
@pytest.mark.parametrize("mesh_device", [pytest.param((8, 4), id="8x4_grid")], indirect=True)
@pytest.mark.parametrize("math_op", [ttnn.ReduceType.Sum])
@pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_line_all_reduce_on_TG_cols_post_commit(
    mesh_device,
    num_devices,
    per_chip_output_shape,
    num_links,
    math_op,
    input_dtype,
    layout,
    buffer_type,
    function_level_defaults,
    replication_factor,
    num_iters=16,
):
    if mesh_device.get_num_devices() != 32:
        pytest.skip("Not TG!")

    run_all_reduce_with_mesh_tensor_along_row(
        mesh_device,
        num_devices,
        per_chip_output_shape,
        num_links,
        math_op,
        input_dtype,
        layout,
        buffer_type,
        function_level_defaults,
        num_iters=num_iters,
        num_all_reduce_instances=replication_factor,
        cluster_axis=0,
    )


@pytest.mark.parametrize(
    "num_devices, num_links, per_chip_output_shape, layout",
    [
        (32, 3, [1, 1, 4096, 50304], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 2048, 50304], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 50304, 4096], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 50304, 2048], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 50304, 1024], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 128000, 4096], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 4096, 128000], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 1024, 50304], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 33, 66], ttnn.TILE_LAYOUT),
        (32, 3, [1, 1, 4094, 50300], ttnn.TILE_LAYOUT),
    ],
)
@pytest.mark.parametrize(
    "input_dtype",
    [
        ttnn.bfloat16,
    ],
)
@pytest.mark.parametrize(
    "buffer_type",
    [
        ttnn.BufferType.DRAM,
    ],
)
@pytest.mark.parametrize("replication_factor", [1])
@pytest.mark.parametrize("math_op", [ttnn.ReduceType.Sum])
@pytest.mark.parametrize("mesh_device", [pytest.param((1, 32), id="1x32_grid")], indirect=True)
@pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_2D}], indirect=True)
def test_line_all_reduce_training(
    mesh_device,
    num_devices,
    per_chip_output_shape,
    num_links,
    math_op,
    input_dtype,
    layout,
    buffer_type,
    function_level_defaults,
    replication_factor,
    num_iters=1,
):
    pytest.skip("Issue 32406: Skipping Failing Test")
    if mesh_device.get_num_devices() != 32:
        pytest.skip("Not TG!")

    run_all_reduce_with_mesh_tensor_along_row(
        mesh_device,
        num_devices,
        per_chip_output_shape,
        num_links,
        math_op,
        input_dtype,
        layout,
        buffer_type,
        function_level_defaults,
        num_iters=num_iters,
        num_all_reduce_instances=replication_factor,
        cluster_axis=1,
    )


@pytest.mark.parametrize("mesh_device", [(4, 8)], indirect=["mesh_device"])
@pytest.mark.parametrize("num_links", [1])
@pytest.mark.parametrize("per_chip_output_shape", [([1, 1, 32, 1280])])
@pytest.mark.parametrize("cluster_axis", [0, 1])
@pytest.mark.parametrize("layout", [ttnn.TILE_LAYOUT])
@pytest.mark.parametrize(
    "input_dtype",
    [ttnn.bfloat16],
)
@pytest.mark.parametrize(
    "memory_config",
    [
        ttnn.MemoryConfig(
            ttnn.TensorMemoryLayout.WIDTH_SHARDED,
            ttnn.BufferType.L1,
            ttnn.ShardSpec(
                ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(4, 4))}),
                [32, 64],
                ttnn.ShardOrientation.ROW_MAJOR,
            ),
        ),
        ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
        ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.L1),
    ],
    ids=["sharded_l1", "interleaved_dram", "interleaved_l1"],
)
@pytest.mark.parametrize("math_op", [ttnn.ReduceType.Sum])
@pytest.mark.parametrize(
    "device_params",
    [{"fabric_config": ttnn.FabricConfig.FABRIC_2D}],
    indirect=True,
    ids=["fabric_2d"],
)
def test_all_reduce_fabric_2d(
    mesh_device,
    per_chip_output_shape,
    cluster_axis,
    num_links,
    math_op,
    input_dtype,
    layout,
    memory_config,
    function_level_defaults,
    device_params,
):
    num_devices = tuple(mesh_device.shape)[cluster_axis]

    run_all_reduce_with_mesh_tensor_along_row(
        mesh_device,
        num_devices,
        per_chip_output_shape,
        num_links,
        math_op,
        input_dtype,
        layout,
        None,
        function_level_defaults,
        memory_config,
        num_iters=1,
        num_all_reduce_instances=1,
        cluster_axis=cluster_axis,
        use_semaphore_free_all_reduce_impl=True,
    )

    if memory_config.is_sharded() == False:
        if device_params["fabric_config"] == ttnn.FabricConfig.FABRIC_2D:
            logger.info(f"Number of program cache entries: {mesh_device.num_program_cache_entries()}")
            assert (
                mesh_device.num_program_cache_entries() == 3
            ), f"Number of program cache entries: {mesh_device.num_program_cache_entries()} but was expecting 3 as we are using fabric 2D, which fallsback to composite all gather + local reduce"
        elif device_params["fabric_config"] == ttnn.FabricConfig.FABRIC_2D:
            logger.info(f"Number of program cache entries: {mesh_device.num_program_cache_entries()}")
            assert (
                mesh_device.num_program_cache_entries() == 2
            ), f"Number of program cache entries: {mesh_device.num_program_cache_entries()} but was expecting 2 as we are using fabric 2D dynamic, which uses reduce scatter + all gather"
